import pymysql
import matplotlib
import matplotlib.pyplot as plt
import pandas as pd
import datetime
import numpy as np

# 初始化数据库连接
mysqlconn = pymysql.connect(host='localhost', port=3306, user='root', password='123456', charset='utf8',
                            database='futuredata_schema')
cursor = mysqlconn.cursor(cursor=pymysql.cursors.DictCursor)  # 设置游标

# 读取合约信息表
df_instrumentInfo = pd.read_sql('select * from futuresinfo_czce', con=mysqlconn)
# print(df_instrumentInfo)
# 只保留所需列
df_instrumentInfo = df_instrumentInfo.loc[:, ['instrumentID', 'instrumentClass', 'listDate', 'expiryDate']]
# print(df_instrumentInfo)
# 将date转化为datetime
df_instrumentInfo['listDate'] = pd.to_datetime(df_instrumentInfo['listDate'])
df_instrumentInfo['expiryDate'] = pd.to_datetime(df_instrumentInfo['expiryDate'])

# 读取合约日线行情数据表,因为数据量大,只读取所需要的列
df_marketData = pd.read_sql('select instrumentID,updateDate,todayClosePrice from futuredaymarketdata_table',
                            con=mysqlconn)
# print(df_marketData)
# 将date转化为datetime
df_marketData['updateDate'] = pd.to_datetime(df_marketData['updateDate'])
# 获取当前日期
g_currentDate = datetime.datetime.now()


# 获取上市天数
def GetListDays(series_listDate):
    days = (g_currentDate - series_listDate).days

    return days


# 获取距离到期天数
def GetExpiryDays(series_expiryDate):
    days = (series_expiryDate - g_currentDate).days
    return days


# 剔除上市天数较短和即将到期的合约
# 计算上市天数
df_instrumentInfo['listDays'] = df_instrumentInfo['listDate'].apply(GetListDays)
# 计算到期天数
df_instrumentInfo['expiryDays'] = df_instrumentInfo['expiryDate'].apply(GetExpiryDays)
# print(df_instrumentInfo)
# 剔除天数不符合条件的合约
minListDays = 40
minExpiryDays = 80
df_instrumentInfo = df_instrumentInfo.drop(df_instrumentInfo[df_instrumentInfo['listDays'] < minListDays].index)
df_instrumentInfo = df_instrumentInfo.drop(df_instrumentInfo[df_instrumentInfo['expiryDays'] < minExpiryDays].index)
# print(df_instrumentInfo)
# 根据合约品种分组,将同品种合约代码按顺序放在一个列表中并排序
# 合约品种-list合约代码
dict_instrumentClassToInstrumentId = {}
for index, row in df_instrumentInfo.iterrows():
    dict_instrumentClassToInstrumentId.setdefault(row['instrumentClass'], []).append(row['instrumentID'])
    dict_instrumentClassToInstrumentId[row['instrumentClass']].sort()
# 检验处理结果
# List = dict_instrumentClassToInstrumentId['AP']
# for value in List:
#     print(value)
# 将总行情数据df按照合约代码分成子df
# 合约代码-行情df
dict_instrumentMarketData = {}
for value in dict_instrumentClassToInstrumentId.values():
    for instrumentID in value:
        # 获取当前合约对应的df
        df_instrumentMarketData = df_marketData[df_marketData['instrumentID'] == instrumentID]
        dict_instrumentMarketData[instrumentID] = df_instrumentMarketData
# print(dict_instrumentMarketData['AP601'])

# 计算价差
# (instrumentIDA-instrumentIDB)-价差df
dict_priceSpread = {}
# 遍历合约品种
for value in dict_instrumentClassToInstrumentId.values():
    listSize = len(value)
    # 两两匹配不重复
    for i in range(0, listSize - 1):
        for j in range(i + 1, listSize):
            instrumentIDA = value[i]
            instrumentIDB = value[j]
            df_A = dict_instrumentMarketData[instrumentIDA]
            df_B = dict_instrumentMarketData[instrumentIDB]
            df_priceSpread = pd.merge(df_A, df_B, how='inner', left_on=['updateDate'], right_on=['updateDate'],
                                      suffixes=['A', 'B'])
            # 计算价差
            df_priceSpread['priceSpread'] = df_priceSpread['todayClosePriceA'] - df_priceSpread['todayClosePriceB']
            # 剔除掉今日收盘价为0时的价差组合（因为冷门合约的存在，此举会导致空dataframe的存在）
            df_priceSpread = df_priceSpread[
                (df_priceSpread['todayClosePriceA'] != 0) & (df_priceSpread['todayClosePriceB'] != 0)]
            df_priceSpread = df_priceSpread.reset_index(drop=True)
            # 只保留所需列
            df_priceSpread = df_priceSpread.loc[:, ['instrumentIDA', 'instrumentIDB', 'updateDate', 'priceSpread']]
            # key形式改为instrumentIDA-instrumentIDB
            dict_priceSpread[instrumentIDA + '-' + instrumentIDB] = df_priceSpread
print(dict_priceSpread['TA601-TA603'])
print(len(dict_priceSpread))

# 绘图
# matplotlib.use('TkAgg')
# df_test = dict_priceSpread['TA601-TA603']
# # plt使用index作为x轴
# df_test.set_index('updateDate', inplace=True)
# plt.plot(df_test['priceSpread'])
# plt.show()

# 计算斜率,回归区间和振幅
df_result = pd.DataFrame()


# 计算斜率的函数
def CalSlope(df):
    list_data = df['priceSpread'].tolist()
    # 输出结果为[斜率,常数]
    slope = np.polyfit(df.index, list_data, deg=1)
    return slope[0]


# 遍历所有价差df，计算斜率（单独计算斜率是为了规避df数据量无法满足polyfit函数的情况)
i = 0
for key, value in dict_priceSpread.items():
    if value.empty or len(value) < 2:
        i += 1
        print(f"跳过空数据或数据不足的键：{key}")
        value['slope'] = np.nan
        continue
    value['slope'] = CalSlope(value)
print(i)

# 遍历所有价差df
for value in dict_priceSpread.values():
    # 计算回归区间
    value['lowerLimit'] = value.groupby(['instrumentIDA', 'instrumentIDB'])['priceSpread'].transform('min')
    value['upperLimit'] = value.groupby(['instrumentIDA', 'instrumentIDB'])['priceSpread'].transform('max')
    # 计算振幅
    value['deviation'] = value['upperLimit'] - value['lowerLimit']
    # 根据合约代码去重
    value.drop_duplicates(['instrumentIDA', 'instrumentIDB'], keep='first', inplace=True)
    # 只保留所需要的列
    value = value.loc[:, ['instrumentIDA', 'instrumentIDB', 'slope', 'lowerLimit', 'upperLimit', 'deviation']]
    df_result = pd.concat([df_result, value], ignore_index=True)
    # df_result = df_result.append(value, ignore_index=True, sort=False) pandas1.4之后不可用

# 查看slope为np.nan的数据有多少条
nan_count_slope = df_result['slope'].isna().sum()
print(nan_count_slope)

# 查看结果有多少条
print(len(df_result))
# df_result.to_csv('df_result.csv')


# 剔除斜率绝对值过大的合约组
maxAbsSlope = 0.5
df_result = df_result.drop(df_result[abs(df_result['slope']) > maxAbsSlope].index)

# 剔除振幅过小的合约组
minDeviation = 20
df_result = df_result.drop(df_result[df_result['deviation'] < minDeviation].index)

# print(df_result)

# 找出斜率绝对值最小的一行
min_slope_row = df_result.loc[df_result['slope'].abs().idxmin()]
# print(min_slope_row)

# 存入csv
df_result.to_csv(r'.\result.csv', index=False, encoding='utf_8_sig')
cursor.close()  # 关闭游标
mysqlconn.close()  # 关闭数据库连接
